from typing import Optional, List, Union
from fastapi import APIRouter
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel
from openai import OpenAI
from config import deepseek, db, apis
import h5py
import os
import json
import shutil
import numpy as np
from method import Log

import time
from datetime import datetime, timedelta

router = APIRouter(default_response_class=JSONResponse)
client = OpenAI(
    api_key=apis.deepseek_key,
    base_url=deepseek.BaseUrl
)

def seconds_to_dhms(seconds):
    td = timedelta(seconds=seconds)
    days = td.days
    hours, remainder = divmod(td.seconds, 3600)
    minutes, seconds = divmod(remainder, 60)
    text = ''
    if days:
        text = f"{days}天"
    if hours:
        text = f"{text}{hours}时"
    if minutes:
        text = f"{text}{minutes}分"
    return text

class Message(BaseModel):
    role: str
    name: str
    content: str

# 正常数据显示
@router.post('/message')
def enter_message(msg: list[Message]):
    try:
        user_msg = []
        messages = []
        messages.append({
            "role": "system",
            "content": deepseek.RoleSystem
        })
        messages.extend(msg)
        response = client.chat.completions.create(
            model=deepseek.model,
            messages=messages,
            max_tokens=2000,
            temperature=1.3,
            stream=False
        )
        # return response.choices[0].message.content
        symbol = [s.content for s in msg if s.name == "symbol"]
        kline = [s.content for s in msg if s.name == "kline"]
        user = [s.content for s in msg if s.name == "user"]
        if kline and kline[0]:
            try:
                kline = json.loads(kline[0])
                formated = "%Y-%m-%d %H:%M"
                if isinstance(kline, list) and kline:
                    start = kline[0].get('OpenTime')
                    end = kline[-1].get('OpenTime')
                    interval = datetime.fromisoformat(kline[1].get('OpenTime')).strftime(formated)
                    # 解析为 datetime 对象
                    start = datetime.fromisoformat(start).strftime(formated)
                    end = datetime.fromisoformat(end).strftime(formated)
                    interval_timestamp = int(datetime.strptime(interval, formated).timestamp())
                    start_timestamp = int(datetime.strptime(start, formated).timestamp())
                    interval_text = seconds_to_dhms(interval_timestamp - start_timestamp)
                    user_msg.append(f"{symbol[0]}({start}~{end}) {interval_text}")
            except Exception as e:
                user_msg.append(f"{symbol[0]} K线数据解析失败")
        if user and user[0]:
            user_msg.append(user[0])
        return {
            'user': user_msg,
            'content': response.choices[0].message.content
        }
    except Exception as e:
        Log.set_log(f"对话失败: {e}", 'error')

@router.post('/message-stream')
def enter_message_stream(msg: List[Message]):
    user_msg = []
    messages = []
    messages.append({
        "role": "system",
        "content": deepseek.RoleSystem
    })
    messages.extend(msg)
    response = client.chat.completions.create(
        model=deepseek.model,
        messages=messages,
        max_tokens=2000,
        temperature=1.3,
        stream=True
    )
    def generate():
        for chunk in response:
            if chunk.choices[0].delta is not None:
                yield f"data: {chunk.model_dump_json()}\n\n"
        yield "data: [DONE]\n\n"
    
    return StreamingResponse(generate(), media_type="text/event-stream")

class Content(BaseModel):
    user: list[str]
    content: str

class Posts(BaseModel):
    symbol: str
    message: Content

@router.post('/message/save')
def save_message(posts: Posts):
    if posts.message:
        # 用时间戳命名 group，避免重复
        timestamp = str(time.time())
        timestamp = timestamp.split('.')[0]

        # 创建一个数据文件（使用 'a' 模式打开文件，避免覆盖已有内容）
        with h5py.File(db.message_data, 'a') as f:
            if posts.symbol in f:
                grp = f[posts.symbol]
            else:
                grp = f.create_group(posts.symbol)
            group = grp.create_group(timestamp)
            # 提取用户消息内容
            group.create_dataset(
                'user',
                data=np.array(posts.message.user, dtype=h5py.string_dtype()),
            )
            group.create_dataset(
                'system',
                data=posts.message.content,
            )
            return True
    return False

class Group(BaseModel):
    group: Optional[str] = ''
    timestamps: Optional[Union[str, int]] = []

@router.post('/message/group')
def group_message(group:Group):
    data_list = []
    if os.path.exists(db.message_data):
        # 读取数据文件（'r' 表示只读）
        with h5py.File(db.message_data, 'r') as f:
            if group.group:
                if group.group in f:
                    grp = f[group.group]
                    members = list(grp.keys())
                    if group.timestamps and len(group.timestamps) >= 2:
                        start = int(group.timestamps[0])
                        end = int(group.timestamps[1])
                    else:
                        start = None  # 或其他默认值
                        end = None
                    def is_even(x):
                        if not start and not end:
                            return x
                        elif not start:
                            return int(x) < end
                        elif not end:
                            return start < int(x)
                        else:
                            return start < int(x) and int(x) < end
                    data_list = list(filter(is_even, members))
            else:
                data_list = list(f.keys())
    return data_list

class Detail(BaseModel):
    group: Optional[str] = ''
    timestamps: Optional[Union[str, int]] = ''
@router.post('/message/detail')
def detail_message(group:Detail):
    key_base = f'{group.group}/{group.timestamps}'
    user_data = []
    system_data = ''
    if os.path.exists(db.message_data):
        # 读取数据文件（'r' 表示只读）
        with h5py.File(db.message_data, 'r') as f:
            # 判断完整 key 路径是否存在
            if f"{key_base}/user" in f and f"{key_base}/system" in f:
                # 用户数据是多个字符串（数组）
                try:
                    user_bytes = f[f"{key_base}/user"][:]
                    user_data = [s.decode('utf-8') for s in user_bytes]
                except Exception as e:
                    print("读取 user 数据失败：", e)
                # 系统数据是单个字符串（标量）
                try:
                    system_bytes = f[f"{key_base}/system"][()]
                    system_data = system_bytes.decode('utf-8')
                except Exception as e:
                    print("读取 system 数据失败：", e)
    return {
        'user': user_data,
        'content': system_data
    }

def compact_h5_file(src_path, dst_path, exclude=[]):
    with h5py.File(src_path, 'r') as src, h5py.File(dst_path, 'w') as dst:
        def copy_func(name, obj):
            if any(name.startswith(e) for e in exclude):
                return
            if isinstance(obj, h5py.Group):
                if name not in dst:
                    dst.create_group(name)
            elif isinstance(obj, h5py.Dataset):
                # 复制数据集到同样路径下
                src.copy(name, dst, name=name)
        src.visititems(copy_func)

@router.get('/message/delete')
def delete_message(path: str):
    try:
        # 示例：删除 path 数据集
        compact_h5_file(db.message_data, db.message_data_copy, exclude=[path])
        # 替换原文件
        shutil.move(db.message_data_copy, db.message_data)
        return {
            'msg': '删除成功',
            'status': True
        }
    except Exception as e:
        return {
            'msg': f'删除失败: {str(e)}',
            'status': False
        }
